
#include "vulkan/VulkanLib.h"

USING_NS_VULKAN

const int Width = 800;
const int Height = 600;

class AppImage : public AppBase
{
public:
    AppImage() : AppBase(Width, Height, "Vulkan Render") {}

    void init() override
    {
        std::vector<Model2DTex::Vertex> vertices{
            {{-1.0f,  1.0f}, {0.0f, 1.0f}, {1.0, 1.0, 1.0}},
            {{-1.0f, -1.0f}, {0.0f, 0.0f}, {1.0, 1.0, 1.0}},
            {{ 1.0f,  1.0f}, {1.0f, 1.0f}, {1.0, 1.0, 1.0}},

            {{ 1.0f,  1.0f}, {1.0f, 1.0f}, {1.0, 1.0, 1.0}},
            {{-1.0f, -1.0f}, {0.0f, 0.0f}, {1.0, 1.0, 1.0}},
            {{ 1.0f, -1.0f}, {1.0f, 0.0f}, {1.0, 1.0, 1.0}},
        };
        model = Model2DTex::create(this->device, vertices);

        texture = Texture::create(this->device, "asset/Lenna.png", true);
        createTextureSampler();

        auto shader = Shader::createByPath("asset/image.vert", "asset/image.frag");
        descriptorSetLayoutAttrib = shader->getDescriptorSetLayout();
        descriptorSetLayout = descriptorSetLayoutAttrib.createVkLayout(device->getDevice());
        this->createUniform();
        this->createDescriptorPool();
        this->createDescriptorSets();

        PipelineConfigInfo pipelineConfig{};
        Pipeline::defaultPipelineConfigInfo(pipelineConfig);
        pipelineConfig.setLayouts.push_back(descriptorSetLayout);
        pipelineConfig.vertexBindings = Model2DTex::Vertex::getBindingDescriptions();
        pipelineConfig.vertexAttributes = Model2DTex::Vertex::getAttributeDescriptions();
        pipelineConfig.renderPass = swapChain->getRenderPass();

        pipeline = Pipeline::create(this->device, shader, pipelineConfig);

        createCommandBuffers();
    }

    void drawFrame() override
    {
        uint32_t imageIndex{};
        auto result = swapChain->acquireNextImage(imageIndex);

        if (result == VK_ERROR_OUT_OF_DATE_KHR)
        {
            recreateSwapChain();
            return;
        }
        else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR)
        {
            throw std::runtime_error("failed to acquire swap chain image");
        }
        updateUniformBuffer(imageIndex);
        recordCommandBuffer(imageIndex);
        result = swapChain->submitCommandBuffers(&commandBuffers[imageIndex], imageIndex);

        if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR)
        {
            recreateSwapChain();
        }
        else if (result != VK_SUCCESS)
        {
            throw std::runtime_error("failed to present swap chain image!");
        }
    }

    void updateUniformBuffer(int i)
    {
        auto& buffer = this->uniformBuffers[i];
        const auto& size = this->window->getSize();
        const auto& imageSize = this->texture->getSize();

        math::Mat4 proj(1.0f);
        math::Mat4 view(1.0f);
        math::Mat4 rotation(1.0f);
        math::Mat4 scale(1.0f);

        proj = math::orgho(-size.x * 0.5f, size.x * 0.5f, -size.y * 0.5f, size.y * 0.5f, -100, 100);
        scale = math::scale({imageSize.x, imageSize.y, 1.0f});

        auto time = (float)TimeSys::getTime();
        constexpr float period = 5.0f;
        rotation = math::rotate(math::pi_2 * std::fmodf(time, period) / period, math::Z_Axis);

        buffer->setField("proj", proj);
        buffer->setField("view", view);
        buffer->setField("model", rotation * scale);

        buffer->flush();
    }

    void cleanup() override
    {
        vkFreeCommandBuffers(device->getDevice(), device->getCommandPool(),
                             static_cast<uint32_t>(commandBuffers.size()), commandBuffers.data());
        commandBuffers.clear();

        model.reset();
        pipeline.reset();
        this->texture.reset();
        this->uniformBuffers.clear();
        
        descriptorSets.clear();

        if (this->textureSampler != nullptr)
        {
            vkDestroySampler(device->getDevice(), textureSampler, nullptr);
            textureSampler = nullptr;
        }

        if (this->descriptorSetLayout != nullptr)
        {
            vkDestroyDescriptorSetLayout(device->getDevice(), descriptorSetLayout, nullptr);
            descriptorSetLayout = nullptr;
        }

        if (this->descriptorPool != nullptr)
        {
            vkDestroyDescriptorPool(device->getDevice(), descriptorPool, nullptr);
            this->descriptorPool = nullptr;
        }
    }

    void createCommandBuffers()
    {
        commandBuffers.resize(swapChain->imageCount());

        VkCommandBufferAllocateInfo alloInfo{};
        alloInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        alloInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        alloInfo.commandPool = device->getCommandPool();
        alloInfo.commandBufferCount = static_cast<uint32_t>(commandBuffers.size());
        if (vkAllocateCommandBuffers(device->getDevice(), &alloInfo, commandBuffers.data()) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to allocate command buffers.");
        }
    }

    void recordCommandBuffer(uint32_t imageIndex)
    {
        VkCommandBufferBeginInfo beginInfo{};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

        if (vkBeginCommandBuffer(commandBuffers[imageIndex], &beginInfo) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to begin recording command buffer!");
        }

        VkRenderPassBeginInfo renderPassInfo{};
        renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
        renderPassInfo.renderPass = swapChain->getRenderPass();
        renderPassInfo.framebuffer = swapChain->getFrameBuffer(imageIndex);
        renderPassInfo.renderArea.offset = { 0, 0 };
        renderPassInfo.renderArea.extent = swapChain->getExtent();

        std::array<VkClearValue, 2> clearValues{};
        clearValues[0].color = { 0.01f, 0.01f, 0.01f, 1.0f };
        clearValues[1].depthStencil = { 1.0f, 0 };
        renderPassInfo.clearValueCount = static_cast<uint32_t>(clearValues.size());
        renderPassInfo.pClearValues = clearValues.data();

        vkCmdBeginRenderPass(commandBuffers[imageIndex], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

        VkViewport viewport{};
        viewport.x = 0.0f;
        viewport.y = 0.0f;
        viewport.width = static_cast<float>(swapChain->getWidth());
        viewport.height = static_cast<float>(swapChain->getHeight());
        viewport.minDepth = 0.0f;
        viewport.maxDepth = 1.0f;
        VkRect2D scissor{ {0, 0}, swapChain->getExtent() };
        vkCmdSetViewport(commandBuffers[imageIndex], 0, 1, &viewport);
        vkCmdSetScissor(commandBuffers[imageIndex], 0, 1, &scissor);

        vkCmdBindPipeline(commandBuffers[imageIndex], VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline->getPipeline());
        vkCmdBindDescriptorSets(commandBuffers[imageIndex], VK_PIPELINE_BIND_POINT_GRAPHICS,
                                pipeline->getPipelineLayout(), 0, 1,
                                &descriptorSets[imageIndex]->getDescriptorSet(), 0, nullptr);

        model->bind(commandBuffers[imageIndex]);
        model->draw(commandBuffers[imageIndex]);

        vkCmdEndRenderPass(commandBuffers[imageIndex]);
        if (vkEndCommandBuffer(commandBuffers[imageIndex]) != VK_SUCCESS)
        {
            throw std::runtime_error("faild to record command buffer");
        }
    }

    void recreateSwapChain()
    {
        vkDeviceWaitIdle(device->getDevice());

        swapChain->recreate();
        if (swapChain->imageCount() != commandBuffers.size())
        {
            vkFreeCommandBuffers(device->getDevice(), device->getCommandPool(),
                                 static_cast<uint32_t>(commandBuffers.size()), commandBuffers.data());
            commandBuffers.clear();

            createCommandBuffers();
        }
    }

    void createUniform()
    {
        this->uniformBuffers.resize(swapChain->imageCount());
        for (int i = 0; i < uniformBuffers.size(); ++i)
        {
            this->uniformBuffers[i] = UniformBuffer::create(device, descriptorSetLayoutAttrib.blocks[0]);
        }
    }

    void createDescriptorPool() {
        std::array<VkDescriptorPoolSize, 2> poolSizes{};
        poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[0].descriptorCount = static_cast<uint32_t>(swapChain->imageCount());
        poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        poolSizes[1].descriptorCount = static_cast<uint32_t>(swapChain->imageCount());

        VkDescriptorPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
        poolInfo.pPoolSizes = poolSizes.data();
        poolInfo.maxSets = static_cast<uint32_t>(swapChain->imageCount());

        if (vkCreateDescriptorPool(device->getDevice(), &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {
            throw std::runtime_error("failed to create descriptor pool!");
        }
    }

    void createDescriptorSets()
    {
        for (size_t i = 0; i < swapChain->imageCount(); i++) {

            auto descriptor = DescriptorSet::create(device, descriptorSetLayout, descriptorPool);

            descriptor->addBuffer(uniformBuffers[i], descriptorSetLayoutAttrib.blocks[0].binding);
            descriptor->addSampler(texture, textureSampler, descriptorSetLayoutAttrib.samples["texSampler"].binging);
            descriptor->build();

            this->descriptorSets.push_back(descriptor);
        }
    }

    void createTextureSampler()
    {
        VkSamplerCreateInfo samplerInfo{};
        samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
        samplerInfo.magFilter = VK_FILTER_LINEAR;
        samplerInfo.minFilter = VK_FILTER_LINEAR;
        samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
        samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
        samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
        samplerInfo.anisotropyEnable = VK_FALSE;
        samplerInfo.maxAnisotropy = 0;
        samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
        samplerInfo.unnormalizedCoordinates = VK_FALSE;
        samplerInfo.compareEnable = VK_FALSE;
        samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;
        samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
        samplerInfo.minLod = 0.0f;
        samplerInfo.maxLod = 0.0f;
        samplerInfo.mipLodBias = 0.0f;

        if (vkCreateSampler(device->getDevice(), &samplerInfo, nullptr, &textureSampler) != VK_SUCCESS) {
            throw std::runtime_error("failed to create texture sampler!");
        }
    }

private:

    Model2DTexRef model;
    PipelineRef  pipeline;
    TextureRef texture;
    DescriptorSetLayout descriptorSetLayoutAttrib;

    VkDescriptorSetLayout descriptorSetLayout{};
    VkDescriptorPool descriptorPool{};
    std::vector<VkCommandBuffer> commandBuffers;
    std::vector<UniformBufferRef> uniformBuffers;
    std::vector<DescriptorSetRef> descriptorSets;
    VkSampler textureSampler{};

};

AppBase* getApplication()
{
    static AppImage app;
    return &app;
}